package cn.edu.sdjzu.knoverse.utils.rag;

import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.input.Prompt;
import dev.langchain4j.model.input.PromptTemplate;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.rag.query.router.QueryRouter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Collection;
import java.util.List;

import static java.util.Arrays.asList;
import static java.util.Collections.emptyList;

/**
 * 自定义查询路由器，支持跳过检索和智能选择检索方式
 */
public class CustomQueryRouter implements QueryRouter {
    private static final Logger log = LoggerFactory.getLogger(CustomQueryRouter.class);
    
    private final ChatLanguageModel chatModel;
    private final ContentRetriever vectorRetriever;
    private final ContentRetriever graphRetriever;
    private final ContentRetriever enhancedRetriever;

    // 用于决定使用哪种检索器的提示模板，包含了是否跳过检索的选项
    private final PromptTemplate retrieverSelectionTemplate = PromptTemplate.from(
            "请分析以下用户查询，并判断应该使用哪种数据源进行检索:\n\n" +
                    "0. 不需要检索: 如果是简单问候、感谢或闲聊，不需要检索知识库\n\n" +
                    "1. 向量数据库: 包含文档内容和详细资料，适合查询具体知识点的详细内容、解释、案例文档等\n\n" +
                    "2. 图数据库: 包含知识图谱，适合查询实体间关系，如课程包含哪些知识点、知识点间的前置关系、案例与知识点的关联等\n\n" +
                    "3. 图增强向量检索: 先从图数据库获取相关实体和关系信息，然后利用这些信息增强向量检索，综合两种数据源的优势\n\n" +
                    "4. 并行检索: 同时使用图数据库和向量数据库，各自独立检索后合并结果\n\n" +
                    "请只回答数字0、1、2、3或4。\n\n" +
                    "用户查询: {{it}}"
    );

    public CustomQueryRouter(ChatLanguageModel chatModel,
                             ContentRetriever vectorRetriever,
                             ContentRetriever graphRetriever,
                             ContentRetriever enhancedRetriever) {
        this.chatModel = chatModel;
        this.vectorRetriever = vectorRetriever;
        this.graphRetriever = graphRetriever;
        this.enhancedRetriever = enhancedRetriever;
    }

    @Override
    public Collection<ContentRetriever> route(Query query) {
        // 直接决定使用哪种检索器或是否跳过检索
        int retrieverType = selectRetrieverType(query.text());
        log.info("检索器选择 ({}): {}", retrieverType, query.text());

        switch (retrieverType) {
            case 0:
                log.info("跳过检索: {}", query.text());
                return emptyList(); // 不需要检索
            case 1:
                return List.of(vectorRetriever); // 只使用向量数据库
            case 2:
                return List.of(graphRetriever);  // 只使用图数据库
            case 3:
                return List.of(enhancedRetriever); // 使用图增强向量检索
            case 4:
            default:
                return asList(vectorRetriever, graphRetriever); // 并行检索
        }
    }

    /**
     * 选择使用哪种检索器或是否跳过检索
     *
     * @return 0=不需要检索, 1=向量数据库, 2=图数据库, 3=图增强向量检索, 4=并行检索
     */
    private int selectRetrieverType(String queryText) {
        Prompt prompt = retrieverSelectionTemplate.apply(queryText);
        AiMessage response = chatModel.chat(prompt.toUserMessage()).aiMessage();
        String answer = response.text().toLowerCase().trim();

        if (answer.contains("0")) {
            return 0; // 不需要检索
        } else if (answer.contains("1")) {
            return 1; // 向量数据库
        } else if (answer.contains("2")) {
            return 2; // 图数据库
        } else if (answer.contains("3")) {
            return 3; // 图增强向量检索
        } else {
            return 4; // 并行检索
        }
    }
} 